Add fake_quantize_per_channel_affine operator with Triton kernel - #5450

Summary

Add a Triton implementation of aten.fake_quantize_per_channel_affine for per-channel affine fake quantization.

The implementation adds the required ATen-compatible behavior:

• Apply an independent scale and zero_point to every channel along the selected axis.
• Compute quantized values with round-half-to-even semantics.
• Clamp quantized values to the requested [quant_min, quant_max] range before dequantization.
• Preserve the input dtype for FP16, FP32, and BF16 tensors.
• Support signed and unsigned quantization ranges.
• Handle arbitrary channel axes, nonzero zero points, non-contiguous inputs, and empty tensors.
• Register the operator in the ATen dispatch table.
• Add a corresponding entry to conf/operators.yaml.
• Add accuracy tests and a standard FlagGems benchmark.

The kernel was generated and iteratively corrected with KernelGen.

Testing

Tested on an NVIDIA H20-3e with PyTorch reference results.

Accuracy tests:

pytest tests/test_fake_quantize_per_channel_affine.py -q

Result:

54 passed

Coverage includes:

• FP16, FP32, and BF16 inputs
• 2D, 3D, and 4D input shapes
• Channel axes from the first through the last dimension
• Unsigned [0, 255] and signed [-128, 127] quantization ranges
• Random nonzero per-channel zero points
• Exact positive and negative half-integer cases for round-half-to-even behavior
• Empty tensors
• Output values and output dtype against the native PyTorch operator

Comprehensive benchmark:

pytest benchmark/test_fake_quantize_per_channel_affine.py \
  --level comprehensive \
  --mode kernel \
  --record json \
  -q --disable-warnings

Result:

1 passed

Pre-commit checks:

python -m pre_commit run --files \
  benchmark/test_fake_quantize_per_channel_affine.py \
  conf/operators.yaml \
  src/flag_gems/__init__.py \
  src/flag_gems/ops/__init__.py \
  src/flag_gems/ops/fake_quantize_per_channel_affine.py \
  tests/test_fake_quantize_per_channel_affine.py

Result:

Check Yaml................Passed
Fix End of Files..........Passed
Trim Trailing Whitespace..Passed
Flake8....................Passed
isort.....................Passed
black.....................Passed
black-jupyter.............Passed

git diff --check upstream/master...HEAD also passed.

Performance

Performance was measured on an NVIDIA H20-3e using the standard FlagGems benchmark framework in kernel mode. All 42 dtype/shape/axis cases completed successfully.

FP16

Average FP16 speedup across these cases: 8.981x.

FP32

Average FP32 speedup across these cases: 8.283x.

BF16

Average BF16 speedup across these cases: 8.998x.

Combined arithmetic average speedup across all 42 cases: 8.754x. Every measured case is faster than the PyTorch reference, with speedups ranging from 7.286x to 10.079x.

Files changed

• src/flag_gems/ops/fake_quantize_per_channel_affine.py

  • Implements the per-channel affine fake-quantization Triton kernel.
  • Maps flattened offsets to the selected channel and loads the corresponding scale and zero point.
  • Implements round-half-to-even, clamping, and dequantization in FP32 while preserving the input/output dtype.

• tests/test_fake_quantize_per_channel_affine.py

  • Adds accuracy tests for supported floating-point dtypes, shapes, axes, quantization ranges, and nonzero zero points.
  • Adds explicit regression coverage for half-to-even rounding and empty tensors.

• benchmark/test_fake_quantize_per_channel_affine.py

  • Adds a standard FlagGems benchmark over multiple shapes, axes, and floating-point dtypes.

• src/flag_gems/ops/__init__.py

  • Exports the implementation.

• src/flag_gems/__init__.py

  • Registers fake_quantize_per_channel_affine in _FULL_CONFIG.

• conf/operators.yaml

  • Adds the operator entry with the KernelGen label and alpha: '5.4' stage.